//
// Created by gaoxiang on 19-5-2.
//

#include <opencv2/opencv.hpp>

#include "myslam/algorithm.h"
#include "myslam/backend.h"
#include "myslam/config.h"
#include "myslam/feature.h"
#include "myslam/frontend.h"
#include "myslam/map.h"
#include "myslam/viewer.h"

namespace myslam {

Frontend::Frontend() {
    num_features_ = Config::Get<int>("num_features");
    num_features_init_ = Config::Get<int>("num_features_init");
    num_features_tracking_ = Config::Get<int>("num_features_tracking");
    num_features_tracking_bad_ = Config::Get<int>("num_features_tracking_bad");
    num_features_needed_for_keyframe_ = Config::Get<int>("num_features_needed_for_keyframe");
    //尝试了FAST、ORB、GFTT，结果还是GFTT效果最好，特征点最均匀
    feature_detector_ = cv::GFTTDetector::create(num_features_, 0.01, 20);
}

//入口函数
bool Frontend::AddFrame(myslam::Frame::Ptr frame) {
    current_frame_ = frame;

    //不管good，bad还是lost，都只能继续追踪。特征点不够就添加关键帧咯
    switch (status_) {
        case FrontendStatus::INITING:
            StereoInit();
            break;
        case FrontendStatus::TRACKING_GOOD:
        case FrontendStatus::TRACKING_BAD:
        case FrontendStatus::LOST:
            Track();
            break;
    }
    //更新上一帧
    last_frame_ = current_frame_;
    return true;
}

bool Frontend::Track() {
    //设置当前帧位姿的初始估计
    //relative_motion_是对上一帧到当前帧运动的估计，它其实就是上上帧到上一帧的运动
    if (last_frame_) {
        current_frame_->SetPose(relative_motion_ * last_frame_->Pose());
    }

    //在当前的左目图像中追踪特征点，更改current_frame_->features_left_
    int num_track_last = TrackLastFrame();
    //利用上一帧的特征点关联的地图点，以及当前帧追踪的特征点，使用BA方法，计算位姿变化
    tracking_inliers_ = EstimateCurrentPoseHW();

    //至少要追踪到50个特征点，才算good
    //如果20个都不到，就算lost
    //但是不管good，bad还是lost，都只能继续追踪。特征点不够就添加关键帧咯
    if (tracking_inliers_ > num_features_tracking_) {
        // tracking good
        status_ = FrontendStatus::TRACKING_GOOD;
        LOG(INFO) << "TRACKING_GOOD";
    } else if (tracking_inliers_ > num_features_tracking_bad_) {
        // tracking bad
        status_ = FrontendStatus::TRACKING_BAD;
        LOG(INFO) << "TRACKING_BAD";
    } else {
        // lost
        status_ = FrontendStatus::LOST;
        LOG(INFO) << "TRACKING_LOST. INSERT KEYFRAME AND KEEP TRACKING";
    }

    //能追踪到的特征点小于80的时候，说明共观已经太少，需要加入新的关键帧，检测新的特征点
    if (tracking_inliers_ < num_features_needed_for_keyframe_)
        InsertKeyframe();

    //更新上一帧到当前帧的运动，作为当前帧到下一帧的运动估计
    relative_motion_ = current_frame_->Pose() * last_frame_->Pose().inverse();

    if (viewer_) viewer_->AddCurrentFrame(current_frame_);
    return true;
}

//新增的关键帧和对应的特征点都会加入到地图中
bool Frontend::InsertKeyframe() {
    // current frame is a new keyframe
    current_frame_->SetKeyFrame();
    map_->InsertKeyFrame(current_frame_);

    LOG(INFO) << "Set frame " << current_frame_->id_ << " as keyframe "
              << current_frame_->keyframe_id_;

    //将mappoint指向关键帧中对应的feature
    SetObservationsForKeyFrame();
    //在关键帧左目图像中检测新的特征点
    DetectFeatures();
    //在右目中寻找匹配
    FindFeaturesInRight();
    //利用左右目视差做特征点三角化
    TriangulateNewPoints();
    //加入新的关键帧后，后端也要更新
    backend_->UpdateMap();

    if (viewer_) viewer_->UpdateMap();

    return true;
}

void Frontend::SetObservationsForKeyFrame() {
    for (auto &feat : current_frame_->features_left_) {
        auto mp = feat->map_point_.lock();
        if (mp) mp->AddObservation(feat);
    }
}

int Frontend::TriangulateNewPoints() {
    std::vector<SE3> poses{camera_left_->pose(), camera_right_->pose()};
    SE3 current_pose_Twc = current_frame_->Pose().inverse();
    int cnt_triangulated_pts = 0;
    for (size_t i = 0; i < current_frame_->features_left_.size(); ++i) {
        if (current_frame_->features_left_[i]->map_point_.expired() &&
            current_frame_->features_right_[i] != nullptr) {
            // 左图的特征点未关联地图点且存在右图匹配点，尝试三角化
            std::vector<Vec3> points{
                camera_left_->pixel2camera(
                    Vec2(current_frame_->features_left_[i]->position_.pt.x,
                         current_frame_->features_left_[i]->position_.pt.y)),
                camera_right_->pixel2camera(
                    Vec2(current_frame_->features_right_[i]->position_.pt.x,
                         current_frame_->features_right_[i]->position_.pt.y))};
            Vec3 pworld = Vec3::Zero();

            if (triangulation(poses, points, pworld) && pworld[2] > 0) {
                auto new_map_point = MapPoint::CreateNewMappoint();
                pworld = current_pose_Twc * pworld;
                new_map_point->SetPos(pworld);
                new_map_point->AddObservation(
                    current_frame_->features_left_[i]);
                new_map_point->AddObservation(
                    current_frame_->features_right_[i]);

                current_frame_->features_left_[i]->map_point_ = new_map_point;
                current_frame_->features_right_[i]->map_point_ = new_map_point;
                map_->InsertMapPoint(new_map_point);
                cnt_triangulated_pts++;
            }
        }
    }
    LOG(INFO) << "new landmarks: " << cnt_triangulated_pts;
    return cnt_triangulated_pts;
}

int Frontend::EstimateCurrentPoseHW(){
    /*
     * 输入什么：
     * current_frame_->Pose()，初始位姿估计
     * camera_left_->K()，左目内参
     * current_frame_->features_left_[i]->map_point_.lock()，地图点坐标
     * current_frame_->features_left_[i]->position_.pt，测量值
     * 做什么：
     * 更改当前帧的位姿估计，current_frame_->SetPose()
     * 更改feat->is_outlier_
     */

    //取mappoint的坐标，以及对应的feature的坐标，避免在多次迭代中反复使用mutex
    std::vector<Vec3> mpts; //地图点坐标
    std::vector<Feature::Ptr> feats; //地图点关联的feature
    for (int i = 0; i < current_frame_->features_left_.size(); ++i) {
        auto mp = current_frame_->features_left_[i]->map_point_.lock();
        if (mp) {
            mpts.emplace_back( mp->Pos() );
            feats.emplace_back(current_frame_->features_left_[i]);
        }
    }
    //取当前帧的相机位姿。避免在迭代中反复取相机位姿，反复使用mutex
    SE3 cam_pose = current_frame_->Pose();

    double cost = 0, lastcost=0;

    double fx = camera_left_->fx_;
    double fy = camera_left_->fy_;
    double cx = camera_left_->cx_;
    double cy = camera_left_->cy_;

    std::vector<double> costs(feats.size()); //存储每个feature的cost，也就是error的二范数，用于判断outlier
    double cost_th = Config::Get<double>("feature_cost_th");
    //进行迭代。迭代结束的情况有两种：
    //dx无解或cost增加：统计outlier，结束迭代，不更新pose
    //dx过小或iter=10：更新pose，统计outlier，结束迭代
    int cnt_outlier = 0;
    int iter;

    //迭代四轮，每轮最多10次
    //每轮后如果有feature发散，将其置为outlier
    for(int ITER=0; ITER<4; ++ITER){
        iter=0;
        while(1){
            Mat66 H = Mat66::Zero();
            Vec6 b = Vec6::Zero();

            cost = 0;
            //遍历所有不是outlier的特征点，计算error和J，累加到H和b
            //第一轮应该所有点都不是outlier
            for (int i = 0; i < feats.size(); i++){
                if(!feats[i]->is_outlier_) {
                    Vec3 pc = cam_pose * mpts[i]; //计算地图点在当前帧相机坐标系中的坐标
                    double X = pc(0, 0);
                    double Y = pc(1, 0);
                    double Z = pc(2, 0);

                    Vec2 e(feats[i]->position_.pt.x - (fx * X / Z + cx), feats[i]->position_.pt.y - (fy * Y / Z + cy));
                    //保存每个feature的cost，用于判断outlier
                    costs[i] = e.squaredNorm();
                    //TODO: 实现huber核，处理outlier
                    //如果某个feature的cost太大，就降低它的权重。算是一种傻瓜式的核函数吧 ～ 囧TZ
                    if(costs[i] > cost_th) e = e*0.1;

                    cost += e.squaredNorm();

                    Mat26 J; //这里的J其实是-J
                    J << -fx / Z, 0, fx * X / Z / Z, fx * X * Y / Z / Z, -fx - fx * X * X / Z / Z, fx * Y / Z,
                            0, -fy / Z, fy * Y / Z / Z, fy + fy * Y * Y / Z / Z, -fy * X * Y / Z / Z, -fy * X / Z;

                    H += J.transpose() * J;
                    b += -J.transpose() * e;
                }
            }
            Eigen::Matrix<double, 1, 6> dx;

            //避免H不正定
            for(int i=0;i<6;i++)
                H(i,i) += 0.1;

            dx = H.ldlt().solve(b);
            //如果dx无解或者cost增长，结束迭代，不更新pose
            if( (std::isnan(dx[0])) || (iter>0&&cost>=lastcost) )
                break;

            //更新pose
            cam_pose = SE3::exp(dx) * cam_pose;
            lastcost =cost;
            iter++;

            //如果iter=10，或者dx已经太小，结束迭代
            if( (iter==10) || (dx.norm()<1e-6) )
                break;
        }
        //根据costs标记outlier
        for (int i = 0; i < feats.size(); i++){
            //std::cout<<"CCCOOOSSSTTT="<<costs[i]<<" IIITTTEEERRR="<<ITER<<std::endl;
            if(costs[i]>cost_th) feats[i]->is_outlier_= true;
            else feats[i]->is_outlier_= false;
        }
    }

    current_frame_->SetPose(cam_pose);

    for(auto feat:feats){
        if(feat->is_outlier_){
            cnt_outlier++;
            feat->map_point_.reset();
            feat->is_outlier_=false;
        }
    }
    LOG(INFO) << "Outlier/Inlier in pose estimating: " << cnt_outlier << "/" << feats.size() - cnt_outlier;
    return feats.size()-cnt_outlier;
}

int Frontend::TrackLastFrame() {
    // use LK flow to estimate points in the left image
    std::vector<cv::Point2f> kps_last, kps_current;
    for (auto &kp : last_frame_->features_left_) {
        if (kp->map_point_.lock()) {
            //如果特征点有对应的地图点，就使用地图点对应的重投影坐标
            auto mp = kp->map_point_.lock();
            auto px = camera_left_->world2pixel(mp->pos_, current_frame_->Pose());
            kps_last.push_back(kp->position_.pt);
            kps_current.push_back(cv::Point2f(px[0], px[1]));
        }
        //否则就使用特征点的像素坐标
        else {
            kps_last.push_back(kp->position_.pt);
            kps_current.push_back(kp->position_.pt);
        }
    }

    std::vector<uchar> status;
    Mat error;
    //在左目相邻两帧中光流追踪点
    //使用特征点在上一帧中的坐标作为初始估计
    cv::calcOpticalFlowPyrLK(
        last_frame_->left_img_, current_frame_->left_img_, kps_last,
        kps_current, status, error, cv::Size(11, 11), 3,
        cv::TermCriteria(cv::TermCriteria::COUNT + cv::TermCriteria::EPS, 30, 0.01),
        cv::OPTFLOW_USE_INITIAL_FLOW);

    int num_good_pts = 0;
    //把成功追踪的特征点加入当前帧
    //如果没成功追踪，也不加空指针
    for (size_t i = 0; i < status.size(); ++i) {
        if (status[i]) {
            cv::KeyPoint kp(kps_current[i], 7);
            Feature::Ptr feature(new Feature(current_frame_, kp));
            feature->map_point_ = last_frame_->features_left_[i]->map_point_;
            current_frame_->features_left_.push_back(feature);
            num_good_pts++;
        }
    }

    LOG(INFO) << "Find " << num_good_pts << " in the last image.";
    return num_good_pts;
}

//find matched features of left image in right image
//init map with features in first left and right image
bool Frontend::StereoInit() {
    //这里DetectFeatures()函数修改了current_frame_->features_left_，将检测到的特征点添加进去
    int num_features_left = DetectFeatures();
    //找到右目照片中和左目匹配的feature
    int num_coor_features = FindFeaturesInRight();
    if (num_coor_features < num_features_init_) {
        return false;
    }

    bool build_map_success = BuildInitMap();
    if (build_map_success) {
        status_ = FrontendStatus::TRACKING_GOOD;
        if (viewer_) {
            viewer_->AddCurrentFrame(current_frame_);
            viewer_->UpdateMap();
        }
        return true;
    }
    return false;
}

int Frontend::DetectFeatures() {
    //为了避免现有的feature被重复检测，搞一个mask
    //原有feature点附近的像素遮挡起来，不检测特征点
    cv::Mat mask(current_frame_->left_img_.size(), CV_8UC1, 255);
    for (auto &feat : current_frame_->features_left_) {
        cv::rectangle(mask, feat->position_.pt - cv::Point2f(10, 10),
                      feat->position_.pt + cv::Point2f(10, 10), 0, CV_FILLED);
    }

    std::vector<cv::KeyPoint> keypoints;
    feature_detector_->detect(current_frame_->left_img_, keypoints, mask);
    int cnt_detected = 0;
    for (auto &kp : keypoints) {
        current_frame_->features_left_.push_back(Feature::Ptr(new Feature(current_frame_, kp)));
        cnt_detected++;
    }

    LOG(INFO) << "Detect " << cnt_detected << " new features";
    return cnt_detected;
}

int Frontend::FindFeaturesInRight() {
    // use LK flow to estimate points in the right image
    std::vector<cv::Point2f> kps_left, kps_right;
    for (auto &kp : current_frame_->features_left_) {
        kps_left.push_back(kp->position_.pt);
        auto mp = kp->map_point_.lock();
        if (mp) {
            // 如果这个特征点对应地图点，那么使用地图点的重投影坐标作为特征点坐标
            auto px = camera_right_->world2pixel(mp->pos_, current_frame_->Pose());
            kps_right.push_back(cv::Point2f(px[0], px[1]));
        } else {
            // use same pixel in left image
            kps_right.push_back(kp->position_.pt);
        }
    }

    std::vector<uchar> status;
    Mat error;
    //这里使用特征点在左目中的坐标作为右目中追踪的初始坐标，使用cv::OPTFLOW_USE_INITIAL_FLOW
    cv::calcOpticalFlowPyrLK(
        current_frame_->left_img_, current_frame_->right_img_, kps_left,
        kps_right, status, error, cv::Size(11, 11), 3,
        cv::TermCriteria(cv::TermCriteria::COUNT + cv::TermCriteria::EPS, 30,
                         0.01),
        cv::OPTFLOW_USE_INITIAL_FLOW);
    //如果在右目中匹配到对应的feature，则将其加入current_frame_->features_right_中，否则就加入一个空指针
    int num_good_pts = 0;
    for (size_t i = 0; i < status.size(); ++i) {
        if (status[i]) {
            cv::KeyPoint kp(kps_right[i], 7);
            Feature::Ptr feat(new Feature(current_frame_, kp));
            feat->is_on_left_image_ = false;
            current_frame_->features_right_.push_back(feat);
            num_good_pts++;
        } else {
            current_frame_->features_right_.push_back(nullptr);
        }
    }
    LOG(INFO) << "Find " << num_good_pts << " in the right image.";
    return num_good_pts;
}

bool Frontend::BuildInitMap() {
    std::vector<SE3> poses{camera_left_->pose(), camera_right_->pose()};
    size_t cnt_init_landmarks = 0;
    for (size_t i = 0; i < current_frame_->features_left_.size(); ++i) {
        if (current_frame_->features_right_[i] == nullptr) continue;
        // create map point from triangulation
        std::vector<Vec3> points{
            camera_left_->pixel2camera(
                Vec2(current_frame_->features_left_[i]->position_.pt.x,
                     current_frame_->features_left_[i]->position_.pt.y)),
            camera_right_->pixel2camera(
                Vec2(current_frame_->features_right_[i]->position_.pt.x,
                     current_frame_->features_right_[i]->position_.pt.y))};
        Vec3 pworld = Vec3::Zero();

        if (triangulation(poses, points, pworld) && pworld[2] > 0) {
            auto new_map_point = MapPoint::CreateNewMappoint();
            new_map_point->SetPos(pworld);
            new_map_point->AddObservation(current_frame_->features_left_[i]);
            new_map_point->AddObservation(current_frame_->features_right_[i]);
            current_frame_->features_left_[i]->map_point_ = new_map_point;
            current_frame_->features_right_[i]->map_point_ = new_map_point;
            cnt_init_landmarks++;
            map_->InsertMapPoint(new_map_point);
        }
    }
    current_frame_->SetKeyFrame();
    map_->InsertKeyFrame(current_frame_);
    backend_->UpdateMap();

    LOG(INFO) << "Initial map created with " << cnt_init_landmarks
              << " map points";

    return true;
}

//reset做什么？
//这一帧的reset由上一帧的lost触发。但lost的情况下上一帧已经检测了新的关键帧，添加了新的特征点
//所以除了继续track，似乎也没有别的更好的选择。总不能直接退出程序吧？
//只能继续track。reset函数目前没有任何用处
bool Frontend::Reset() {
    LOG(INFO) << "reset, and then? just keep tracking";
    return false;
}

}  // namespace myslam
